import numpy as np
from scipy.linalg import sqrtm, inv, det
import matplotlib.pyplot as plt
from scipy.io import loadmat

# --- 定义模型函数 ---
def fx_cv(x):
    """匀速(CV)模型状态转移函数"""
    # x = [x, vx, y, vy, ax, ay]^T (虽然CV模型中ax, ay为0或常数，但状态向量统一)
    F = np.array([[1, T, 0, 0, 0, 0],
                  [0, 1, 0, 0, 0, 0],
                  [0, 0, 1, T, 0, 0],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])
    return F @ x

def fx_ca(x):
    """匀加速(CA)模型状态转移函数"""
    # x = [x, vx, y, vy, ax, ay]^T
    F = np.array([[1, T, 0, 0, 0.5*T**2, 0],
                  [0, 1, 0, 0, T, 0],
                  [0, 0, 1, T, 0, 0.5*T**2],
                  [0, 0, 0, 1, 0, T],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])
    return F @ x

def h(x):
    """观测函数 (从笛卡尔坐标转换为极坐标: 距离, 角度)"""
    px, _, py, _, _, _ = x
    r = np.sqrt(px**2 + py**2)
    if r < 1e-6: # 防止除以零
        r = 1e-6
    theta = np.arctan2(py, px)
    return np.array([r, theta])

# --- 主程序 ---
def run_imm_ukf():
    # --- 1. 数据加载 ---
    try:
        # 假设 .mat 文件中的变量名与 MATLAB 工作区中的变量名一致
        measure_data = loadmat('Measure.mat')
        real_data = loadmat('Real.mat')

        # 提取量测数据 TimeRPhi (应为 2x501)
        # 注意：loadmat 加载的变量通常是二维数组，即使在MATLAB中是向量
        Z_full = measure_data['TimeRPhi'] # 假设变量名就是 TimeRPhi
        Z = Z_full[1:3, :] # 取第2、3行作为量测值 (R, Phi)

        # 提取真实值数据 TXYdXdYAxAyTPhi (应为 7x501)
        Xr_full = real_data['TXYdXdYAxAyTPhi'] # 假设变量名就是 TXYdXdYAxAyTPhi
        # Xr_full 结构: [T, X, Y, dX, dY, Ax, Ay, Phi] (根据注释推测)
        # 我们需要 [X, Y, dX, dY, Ax, Ay] 对应状态 [x, vx, y, vy, ax, ay]
        Xr = Xr_full[1:7, :] # 取第2到7行作为真实状态

    except FileNotFoundError as e:
        print(f"错误：找不到.mat文件。请确保 Measure.mat 和 Real.mat 在当前目录下。{e}")
        return
    except KeyError as e:
        print(f"错误：在.mat文件中找不到指定变量。请检查变量名是否正确。{e}")
        return

    num_steps = Z.shape[1] # 501

    # --- 2. 参数初始化 ---
    T = 0.1 # 采样周期
    # 马尔可夫转移概率矩阵 (从列到行: P(当前模型|上一模型))
    pi = np.array([[0.8, 0.2],
                   [0.1, 0.9]])

    # UKF 参数
    alpha = 0.001
    beta = 2
    n = 6 # 状态维度
    kappa = (alpha**2 - 1) * n
    lambda_ukf = alpha**2 * (n + kappa) - n

    # 权重计算
    Wm0 = lambda_ukf / (n + lambda_ukf)
    Wc0 = lambda_ukf / (n + lambda_ukf) + (1 - alpha**2 + beta)
    Wm_other = 1 / (2 * (n + lambda_ukf))
    Wc_other = 1 / (2 * (n + lambda_ukf))

    # 初始状态和协方差 (两个模型)
    X1 = np.array([10000, 0, 0, 0, 0, 0], dtype=float) # CV模型初始状态
    P1 = np.diag([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])      # CV模型初始协方差

    X2 = np.array([10000, 0, 0, 0, 0, 0], dtype=float) # CA模型初始状态
    P2 = np.diag([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])      # CA模型初始协方差

    # 噪声协方差
    R = np.diag([100, (0.1 / 57.3)**2]) # 量测噪声协方差 (R, Phi)

    # Q1 (CV模型过程噪声) - 在原代码中设为0
    Q1 = np.zeros((n, n))

    # Q2 (CA模型过程噪声)
    # Ta2 是噪声驱动矩阵 G
    Ta2 = np.array([[0.5*T**2, 0],
                    [T, 0],
                    [0, 0.5*T**2],
                    [0, T],
                    [1, 0],
                    [0, 1]])
    sigma2 = 7
    Q2 = Ta2 @ Ta2.T * sigma2**2

    # 初始模型概率
    mu_cv = 0.8 # 初始时刻CV模型概率
    mu_ca = 0.2 # 初始时刻CA模型概率 (注意：原代码注释与值不一致，这里按值和逻辑推)

    # 存储结果
    Xe = np.zeros((n, num_steps)) # 融合估计
    Xe1 = np.zeros((n, num_steps)) # CV模型估计
    Xe2 = np.zeros((n, num_steps)) # CA模型估计
    mu_probs = np.zeros((2, num_steps)) # 存储模型概率

    # --- 3. 主循环 ---
    for k in range(num_steps):
        # print(f"Processing step {k+1}/{num_steps}") # 可选：进度提示

        # --- 第一步：交互混合 (Model Condition Re-initialization) ---
        # 1. 计算混合概率 mu_ij = P(M_k = j | M_{k-1} = i)
        c1 = pi[0, 0] * mu_cv + pi[1, 0] * mu_ca
        c2 = pi[0, 1] * mu_cv + pi[1, 1] * mu_ca

        # 2. 计算混合后的状态和协方差
        # CV模型 (j=0) 的混合
        mu11 = pi[0, 0] * mu_cv / c1 if c1 > 0 else 0
        mu21 = pi[1, 0] * mu_ca / c1 if c1 > 0 else 0
        X1_mix = X1 * mu11 + X2 * mu21
        # P1_mix = mu11 * (P1 + (X1 - X1_mix) @ (X1 - X1_mix).T) + \
        #          mu21 * (P2 + (X2 - X1_mix) @ (X2 - X1_mix).T)
        # 优化：避免重复计算 X1-X1_mix
        dx1 = X1 - X1_mix
        dx2 = X2 - X1_mix
        P1_mix = mu11 * (P1 + np.outer(dx1, dx1)) + \
                 mu21 * (P2 + np.outer(dx2, dx2))

        # CA模型 (j=1) 的混合
        mu12 = pi[0, 1] * mu_cv / c2 if c2 > 0 else 0
        mu22 = pi[1, 1] * mu_ca / c2 if c2 > 0 else 0
        X2_mix = X1 * mu12 + X2 * mu22
        # P2_mix = mu12 * (P1 + (X1 - X2_mix) @ (X1 - X2_mix).T) + \
        #          mu22 * (P2 + (X2 - X2_mix) @ (X2 - X2_mix).T)
        dx1 = X1 - X2_mix
        dx2 = X2 - X2_mix
        P2_mix = mu12 * (P1 + np.outer(dx1, dx1)) + \
                 mu22 * (P2 + np.outer(dx2, dx2))


        # --- 第二步：模型条件滤波 (Model-Conditioned Filtering) ---
        # --- 对 CV 模型进行 UKF ---
        # 1. 生成 Sigma 点
        try:
            A1 = sqrtm(P1_mix)
            if np.iscomplexobj(A1):
                 # 如果矩阵不是正定的，sqrtm可能返回复数
                 # 可以取实部，或使用其他方法如 cholesky (如果正定)
                 print(f"Warning: sqrtm returned complex for P1_mix at step {k+1}. Taking real part.")
                 A1 = np.real(A1)
            A1 = A1.T # 与MATLAB 'A1=A1'' 一致
        except np.linalg.LinAlgError:
             print(f"Error computing sqrtm for P1_mix at step {k+1}. Using zero matrix.")
             A1 = np.zeros_like(P1_mix)

        # 生成 2n+1 个 Sigma 点 (这里 n=6, 所以是 13 个点)
        sigma_points_1 = np.zeros((n, 2 * n + 1))
        sigma_points_1[:, 0] = X1_mix
        for i in range(n):
            sigma_points_1[:, i + 1] = X1_mix + np.sqrt(n + lambda_ukf) * A1[:, i]
            sigma_points_1[:, i + 1 + n] = X1_mix - np.sqrt(n + lambda_ukf) * A1[:, i]

        # 2. 时间更新 (预测)
        # 通过状态转移函数传递 Sigma 点
        predicted_sigma_points_1 = np.zeros_like(sigma_points_1)
        for i in range(2 * n + 1):
            predicted_sigma_points_1[:, i] = fx_cv(sigma_points_1[:, i])

        # 预测状态均值 X1_k|k-1
        X1_pred = Wm0 * predicted_sigma_points_1[:, 0]
        for i in range(1, 2 * n + 1):
            X1_pred += Wm_other * predicted_sigma_points_1[:, i]

        # 预测状态协方差 P1_k|k-1
        P1_pred = Wc0 * np.outer(predicted_sigma_points_1[:, 0] - X1_pred,
                                 predicted_sigma_points_1[:, 0] - X1_pred)
        for i in range(1, 2 * n + 1):
            diff = predicted_sigma_points_1[:, i] - X1_pred
            P1_pred += Wc_other * np.outer(diff, diff)
        P1_pred += Q1 # 加上过程噪声

        # 3. 观测更新 (更新)
        # 通过观测函数传递预测的 Sigma 点
        predicted_meas_points_1 = np.zeros((2, 2 * n + 1))
        for i in range(2 * n + 1):
            predicted_meas_points_1[:, i] = h(predicted_sigma_points_1[:, i])

        # 预测观测均值 Z1_k|k-1
        Z1_pred = Wm0 * predicted_meas_points_1[:, 0]
        for i in range(1, 2 * n + 1):
            Z1_pred += Wm_other * predicted_meas_points_1[:, i]

        # 预测观测协方差 Pzz1
        Pzz1 = Wc0 * np.outer(predicted_meas_points_1[:, 0] - Z1_pred,
                              predicted_meas_points_1[:, 0] - Z1_pred)
        for i in range(1, 2 * n + 1):
            diff = predicted_meas_points_1[:, i] - Z1_pred
            Pzz1 += Wc_other * np.outer(diff, diff)
        Pzz1 += R # 加上观测噪声

        # 状态-观测互协方差 Pxz1
        Pxz1 = Wc0 * np.outer(predicted_sigma_points_1[:, 0] - X1_pred,
                              predicted_meas_points_1[:, 0] - Z1_pred)
        for i in range(1, 2 * n + 1):
            diff_x = predicted_sigma_points_1[:, i] - X1_pred
            diff_z = predicted_meas_points_1[:, i] - Z1_pred
            Pxz1 += Wc_other * np.outer(diff_x, diff_z)

        # UKF 增益 K1
        # K1 = Pxz1 * inv(Pzz1) # 使用 inv 可能数值不稳定
        try:
            K1 = Pxz1 @ inv(Pzz1)
        except np.linalg.LinAlgError:
            print(f"Error inverting Pzz1 at step {k+1} for model 1. Skipping update.")
            K1 = np.zeros_like(Pxz1)

        # 更新状态和协方差
        innov1 = Z[:, k] - Z1_pred # 新息
        X1_updated = X1_pred + K1 @ innov1
        P1_updated = P1_pred - K1 @ Pzz1 @ K1.T

        # --- 对 CA 模型进行 UKF --- (与 CV 模型类似)
        try:
            A2 = sqrtm(P2_mix)
            if np.iscomplexobj(A2):
                 print(f"Warning: sqrtm returned complex for P2_mix at step {k+1}. Taking real part.")
                 A2 = np.real(A2)
            A2 = A2.T
        except np.linalg.LinAlgError:
             print(f"Error computing sqrtm for P2_mix at step {k+1}. Using zero matrix.")
             A2 = np.zeros_like(P2_mix)

        sigma_points_2 = np.zeros((n, 2 * n + 1))
        sigma_points_2[:, 0] = X2_mix
        for i in range(n):
            sigma_points_2[:, i + 1] = X2_mix + np.sqrt(n + lambda_ukf) * A2[:, i]
            sigma_points_2[:, i + 1 + n] = X2_mix - np.sqrt(n + lambda_ukf) * A2[:, i]

        predicted_sigma_points_2 = np.zeros_like(sigma_points_2)
        for i in range(2 * n + 1):
            predicted_sigma_points_2[:, i] = fx_ca(sigma_points_2[:, i])

        X2_pred = Wm0 * predicted_sigma_points_2[:, 0]
        for i in range(1, 2 * n + 1):
            X2_pred += Wm_other * predicted_sigma_points_2[:, i]

        P2_pred = Wc0 * np.outer(predicted_sigma_points_2[:, 0] - X2_pred,
                                 predicted_sigma_points_2[:, 0] - X2_pred)
        for i in range(1, 2 * n + 1):
            diff = predicted_sigma_points_2[:, i] - X2_pred
            P2_pred += Wc_other * np.outer(diff, diff)
        P2_pred += Q2

        predicted_meas_points_2 = np.zeros((2, 2 * n + 1))
        for i in range(2 * n + 1):
            predicted_meas_points_2[:, i] = h(predicted_sigma_points_2[:, i])

        Z2_pred = Wm0 * predicted_meas_points_2[:, 0]
        for i in range(1, 2 * n + 1):
            Z2_pred += Wm_other * predicted_meas_points_2[:, i]

        Pzz2 = Wc0 * np.outer(predicted_meas_points_2[:, 0] - Z2_pred,
                              predicted_meas_points_2[:, 0] - Z2_pred)
        for i in range(1, 2 * n + 1):
            diff = predicted_meas_points_2[:, i] - Z2_pred
            Pzz2 += Wc_other * np.outer(diff, diff)
        Pzz2 += R

        Pxz2 = Wc0 * np.outer(predicted_sigma_points_2[:, 0] - X2_pred,
                              predicted_meas_points_2[:, 0] - Z2_pred)
        for i in range(1, 2 * n + 1):
            diff_x = predicted_sigma_points_2[:, i] - X2_pred
            diff_z = predicted_meas_points_2[:, i] - Z2_pred
            Pxz2 += Wc_other * np.outer(diff_x, diff_z)

        try:
            K2 = Pxz2 @ inv(Pzz2)
        except np.linalg.LinAlgError:
            print(f"Error inverting Pzz2 at step {k+1} for model 2. Skipping update.")
            K2 = np.zeros_like(Pxz2)

        innov2 = Z[:, k] - Z2_pred
        X2_updated = X2_pred + K2 @ innov2
        P2_updated = P2_pred - K2 @ Pzz2 @ K2.T

        # --- 第三步：模型概率更新 (Model Probability Update) ---
        # 计算似然函数 (模型匹配程度)
        # likeli = exp(-0.5 * innov' * inv(Pzz) * innov) / sqrt(det(2*pi*Pzz))
        # 取对数计算更稳定: log_likeli = -0.5 * (innov' * inv(Pzz) * innov + log(det(2*pi*Pzz)))
        # 原代码直接计算指数部分，这里也直接计算，但增加稳定性检查
        def compute_likelihood(innov, Pzz):
             try:
                 # 避免 Pzz 奇异
                 det_Pzz = det(Pzz)
                 if det_Pzz <= 0:
                     print("Warning: Non-positive determinant in likelihood calculation.")
                     return 1e-100 # 返回一个极小的似然
                 norm_factor = np.sqrt(det_Pzz * (2 * np.pi)**len(innov))
                 if norm_factor == 0:
                     return 1e-100
                 # innov_T_Pzz_inv_innov = innov.T @ inv(Pzz) @ innov
                 # 使用 solve 更稳定
                 innov_T_Pzz_inv_innov = innov @ np.linalg.solve(Pzz, innov)
                 likelihood = np.exp(-0.5 * innov_T_Pzz_inv_innov) / norm_factor
                 return likelihood if likelihood > 0 else 1e-100
             except (np.linalg.LinAlgError, FloatingPointError):
                 print("Numerical error in likelihood computation. Returning small value.")
                 return 1e-100

        likelihood1 = compute_likelihood(innov1, Pzz1)
        likelihood2 = compute_likelihood(innov2, Pzz2)

        # 更新模型概率
        # c = likelihood1 * c1 + likelihood2 * c2 # 归一化常数
        # 为避免数值下溢，使用对数域计算
        log_c1_term = np.log(likelihood1) + np.log(c1) if c1 > 0 and likelihood1 > 0 else -np.inf
        log_c2_term = np.log(likelihood2) + np.log(c2) if c2 > 0 and likelihood2 > 0 else -np.inf
        # log_sum_exp
        max_log = max(log_c1_term, log_c2_term)
        if max_log == -np.inf:
            # 所有项都是0概率，分配均匀概率或保持原状（这里简单处理）
            mu_cv_new, mu_ca_new = mu_cv, mu_ca
            print(f"Warning: All likelihoods zero at step {k+1}. Keeping previous probabilities.")
        else:
            log_c = max_log + np.log(np.exp(log_c1_term - max_log) + np.exp(log_c2_term - max_log))
            mu_cv_new = np.exp(log_c1_term - log_c) if log_c > -np.inf else 0
            mu_ca_new = np.exp(log_c2_term - log_c) if log_c > -np.inf else 0

        mu_cv, mu_ca = mu_cv_new, mu_ca_new

        # --- 第四步：估计融合 (Estimate Fusion) ---
        X_fused = X1_updated * mu_cv + X2_updated * mu_ca

        # --- 存储结果 ---
        Xe1[:, k] = X1_updated
        Xe2[:, k] = X2_updated
        Xe[:, k] = X_fused
        mu_probs[0, k] = mu_cv
        mu_probs[1, k] = mu_ca

        # --- 更新下一时刻的先验 ---
        X1, P1 = X1_updated, P1_updated
        X2, P2 = X2_updated, P2_updated

    # --- 5. 结果可视化 ---
    time_steps = np.arange(num_steps)

    plt.figure(figsize=(15, 10))
    states = ['X Position', 'X Velocity', 'Y Position', 'Y Velocity', 'X Acceleration', 'Y Acceleration']
    for i in range(n):
        plt.subplot(3, 2, i+1)
        plt.plot(time_steps, Xr[i, :] - Xe[i, :], label='Fused Error')
        plt.plot(time_steps, Xr[i, :] - Xe1[i, :], '--', label='CV Model Error', alpha=0.7)
        plt.plot(time_steps, Xr[i, :] - Xe2[i, :], '-.', label='CA Model Error', alpha=0.7)
        plt.title(f'Error in {states[i]}')
        plt.xlabel('Time Step')
        plt.ylabel('Error')
        plt.legend()
        plt.grid(True)
    plt.tight_layout()
    plt.show()

    # 绘制模型概率
    plt.figure(figsize=(10, 4))
    plt.plot(time_steps, mu_probs[0, :], label='P(CV)')
    plt.plot(time_steps, mu_probs[1, :], label='P(CA)')
    plt.title('Model Probabilities')
    plt.xlabel('Time Step')
    plt.ylabel('Probability')
    plt.legend()
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    run_imm_ukf()




